#***************************************************************************
#*                                                                         *
#*   Copyright (c) 2018 Yorik van Havre <yorik@uncreated.net>              *
#*                                                                         *
#*   This program is free software; you can redistribute it and/or modify  *
#*   it under the terms of the GNU Lesser General Public License (LGPL)    *
#*   as published by the Free Software Foundation; either version 2 of     *
#*   the License, or (at your option) any later version.                   *
#*   for detail see the LICENCE text file.                                 *
#*                                                                         *
#*   This program is distributed in the hope that it will be useful,       *
#*   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
#*   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
#*   GNU Library General Public License for more details.                  *
#*                                                                         *
#*   You should have received a copy of the GNU Library General Public     *
#*   License along with this program; if not, write to the Free Software   *
#*   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  *
#*   USA                                                                   *
#*                                                                         *
#***************************************************************************

from __future__ import print_function

"""This module contains FreeCAD commands for the BIM workbench"""

import os
import FreeCAD
from BimTranslateUtils import *
import BimIfcElements

class BIM_Classification:


    def GetResources(self):

        return {'Pixmap'  : os.path.join(os.path.dirname(__file__),"icons","BIM_Classification.svg"),
                'MenuText': QT_TRANSLATE_NOOP("BIM_Classification", "Manage classification..."),
                'ToolTip' : QT_TRANSLATE_NOOP("BIM_Classification", "Manage how the different materials of this documents use classification systems")}

    def IsActive(self):

        if FreeCAD.ActiveDocument:
            return True
        else:
            return False

    def Activated(self):

        import FreeCADGui,Draft
        from PySide import QtCore,QtGui

        # init checks
        if not hasattr(self,"Classes"):
            self.Classes = {}
        self.isEditing = None

        # load the form and set the tree model up
        self.form = FreeCADGui.PySideUic.loadUi(os.path.join(os.path.dirname(__file__),"dialogClassification.ui"))
        self.form.setWindowIcon(QtGui.QIcon(os.path.join(os.path.dirname(__file__),"icons","BIM_Classification.svg")))

        # restore saved values
        p = FreeCAD.ParamGet("User parameter:BaseApp/Preferences/Mod/BIM")
        w = p.GetInt("BimClassificationDialogWidth",629)
        h = p.GetInt("BimClassificationDialogHeight",516)
        self.form.resize(w,h)

        # add modified search box from bimmaterial
        import BimMaterial
        searchBox = BimMaterial.MatLineEdit(self.form)
        searchBox.setPlaceholderText(translate("BIM","Search..."))
        searchBox.setToolTip(translate("BIM","Searches classes"))
        self.form.search = searchBox
        self.form.horizontalLayout_2.addWidget(searchBox)

        # set help line
        self.form.labelDownload.setText(self.form.labelDownload.text().replace("%s",os.path.join(FreeCAD.getUserAppDataDir(),"BIM","Classification")))

        # hide materials list if we are editing a particular object
        if len(FreeCADGui.Selection.getSelection()) == 1:
            self.isEditing = FreeCADGui.Selection.getSelection()[0]
            if hasattr(self.isEditing,"StandardCode"):
                self.form.groupMaterials.hide()
                self.form.buttonApply.hide()
                self.form.buttonRename.hide()
                self.form.setWindowTitle(translate("BIM","Editing")+" "+self.isEditing.Label)

        # fill materials list
        self.objectslist = {}
        self.matlist = {}
        self.labellist = {}
        for obj in FreeCAD.ActiveDocument.Objects:
            if "StandardCode" in obj.PropertiesList:
                if Draft.getType(obj) in ["Material","MultiMaterial"]:
                    self.matlist[obj.Name] = obj.StandardCode
                else:
                    self.objectslist[obj.Name] = obj.StandardCode
                self.labellist[obj.Name] = obj.Label

        # fill objects list
        if not self.isEditing:
            self.updateObjects()

        # fill available classifications
        p = FreeCAD.ParamGet("User parameter:BaseApp/Preferences/Mod/Arch").GetString("DefaultClassificationSystem","")
        presetdir = os.path.join(FreeCAD.getUserAppDataDir(),"BIM","Classification")
        if os.path.isdir(presetdir):
            presets = []
            for f in os.listdir(presetdir):
                if f.lower().endswith(".xml") or f.lower().endswith(".ifc"):
                    n = os.path.splitext(f)[0]
                    if not n in presets:
                        presets.append(n)
                        self.form.comboSystem.addItem(n)
                        if n == p:
                            self.form.comboSystem.setCurrentIndex(self.form.comboSystem.count()-1)

        # connect signals
        QtCore.QObject.connect(self.form.comboSystem, QtCore.SIGNAL("currentIndexChanged(int)"), self.updateClasses)
        QtCore.QObject.connect(self.form.buttonApply, QtCore.SIGNAL("clicked()"), self.apply)
        QtCore.QObject.connect(self.form.buttonRename, QtCore.SIGNAL("clicked()"), self.rename)
        QtCore.QObject.connect(self.form.search, QtCore.SIGNAL("textEdited(QString)"), self.updateClasses)
        QtCore.QObject.connect(self.form.buttonBox, QtCore.SIGNAL("accepted()"), self.accept)
        QtCore.QObject.connect(self.form.groupMode, QtCore.SIGNAL("currentIndexChanged(int)"), self.updateObjects)
        QtCore.QObject.connect(self.form.treeClass, QtCore.SIGNAL("itemDoubleClicked(QTreeWidgetItem*,int)"), self.apply)
        QtCore.QObject.connect(self.form.search,QtCore.SIGNAL("up()"),self.onUpArrow)
        QtCore.QObject.connect(self.form.search,QtCore.SIGNAL("down()"),self.onDownArrow)

        # center the dialog over FreeCAD window
        mw = FreeCADGui.getMainWindow()
        self.form.move(mw.frameGeometry().topLeft() + mw.rect().center() - self.form.rect().center())

        self.updateClasses()
        self.form.show()

        self.form.search.setFocus()

    def updateObjects(self,idx=None):

        # store current state of tree into self.objectslist before redrawing

        for row in range(self.form.treeObjects.topLevelItemCount()):
            child = self.form.treeObjects.topLevelItem(row)
            if child.toolTip(0):
                if child.toolTip(0) in self.objectslist:
                    self.objectslist[child.toolTip(0)] = child.text(1)
                elif child.toolTip(0) in self.matlist:
                    self.matlist[child.toolTip(0)] = child.text(1)
                self.labellist[child.toolTip(0)] = child.text(0)
            for childrow in range(child.childCount()):
                grandchild = child.child(childrow)
                if grandchild.toolTip(0):
                    if grandchild.toolTip(0) in self.objectslist:
                        self.objectslist[grandchild.toolTip(0)] = grandchild.text(1)
                    elif grandchild.toolTip(0) in self.matlist:
                        self.matlist[grandchild.toolTip(0)] = grandchild.text(1)
                    self.labellist[grandchild.toolTip(0)] = grandchild.text(0)

        self.form.treeObjects.clear()

        if self.form.groupMode.currentIndex() == 1:
            # group by type
            self.updateByType()
        elif self.form.groupMode.currentIndex() == 2:
            # group by material
            self.updateByMaterial()
        elif self.form.groupMode.currentIndex() == 3:
            # group by model structure
            self.updateByTree()
        else:
            # group alphabetically
            self.updateDefault()

        # resize columns - no resizeSection in pyside2
        #self.form.treeObjects.header().resizeSection(0,int(self.form.treeObjects.width()/2))
        #self.form.treeObjects.header().resizeSection(1,int(self.form.treeObjects.width()/2))

    def updateByType(self):

        from PySide import QtCore,QtGui
        import Draft
        groups = {}
        for name in self.objectslist.keys():
            obj = FreeCAD.ActiveDocument.getObject(name)
            if obj and hasattr(obj,"IfcType"):
                groups.setdefault(obj.IfcType,[]).append(name)
            elif obj and hasattr(obj,"IfcRole"):
                groups.setdefault(obj.IfcRole,[]).append(name)
            else:
                groups.setdefault("Undefined",[]).append(name)
        groups["Materials"] = self.matlist.keys()
        d = self.objectslist.copy()
        d.update(self.matlist)

        for group in groups.keys():
            mit = QtGui.QTreeWidgetItem([group,""])
            self.form.treeObjects.addTopLevelItem(mit)
            for name in groups[group]:
                obj = FreeCAD.ActiveDocument.getObject(name)
                if obj:
                    if (not self.form.onlyVisible.isChecked()) or obj.ViewObject.isVisible() or (Draft.getType(obj) in ["Material","MultiMaterial"]):
                        it = QtGui.QTreeWidgetItem([self.labellist[name],d[name]])
                        it.setIcon(0,BimIfcElements.getIcon(obj))
                        it.setToolTip(0,name)
                        mit.addChild(it)
            mit.sortChildren(0,QtCore.Qt.AscendingOrder)
        self.form.treeObjects.expandAll()
        #self.spanTopLevels()

    def updateByMaterial(self):

        from PySide import QtCore,QtGui
        groups = {}
        claimed = []
        for name in self.matlist.keys():
            mat = FreeCAD.ActiveDocument.getObject(name)
            if mat:
                children = [par.Name for par in mat.InList if par.Name in self.objectslist.keys()]
                groups[name] = children
                claimed.extend(children)
        groups["Undefined"] = [o for o in self.objectslist.keys() if not o in claimed]

        for group in groups.keys():
            matobj = FreeCAD.ActiveDocument.getObject(group)
            if matobj:
                mit = QtGui.QTreeWidgetItem([self.labellist[group],self.matlist[group]])
                mit.setIcon(0,BimIfcElements.getIcon(matobj))
                mit.setToolTip(0,group)
            else:
                mit = QtGui.QTreeWidgetItem(["Undefined",""])
            self.form.treeObjects.addTopLevelItem(mit)
            for name in groups[group]:
                obj = FreeCAD.ActiveDocument.getObject(name)
                if obj:
                    if (not self.form.onlyVisible.isChecked()) or obj.ViewObject.isVisible():
                        it = QtGui.QTreeWidgetItem([self.labellist[name],self.objectslist[name]])
                        it.setIcon(0,BimIfcElements.getIcon(obj))
                        it.setToolTip(0,name)
                        mit.addChild(it)
            mit.sortChildren(0,QtCore.Qt.AscendingOrder)
        self.form.treeObjects.expandAll()
        #self.spanTopLevels()

    def updateByTree(self):

        from PySide import QtCore,QtGui
        # order by hierarchy
        def istop(obj):
            for parent in obj.InList:
                if parent.Name in self.objectslist.keys():
                    return False
            return True

        rel = []
        deps = []
        for name in self.objectslist.keys():
            obj = FreeCAD.ActiveDocument.getObject(name)
            if obj:
                if istop(obj):
                    rel.append(obj)
                else:
                    deps.append(obj)
        pa = 1
        while deps:
            for obj in rel:
                for child in obj.OutList:
                    if child in deps:
                        rel.append(child)
                        deps.remove(child)
            pa += 1
            if pa == 10: # max 10 hierarchy levels, okay? Let's keep civilised
                rel.extend(deps)
                break

        done = {}
        # materials first
        mit = QtGui.QTreeWidgetItem(["Materials",""])
        self.form.treeObjects.addTopLevelItem(mit)
        for name,code in self.matlist.items():
            obj = FreeCAD.ActiveDocument.getObject(name)
            if obj:
                it = QtGui.QTreeWidgetItem([self.labellist[name],code])
                it.setIcon(0,BimIfcElements.getIcon(obj))
                it.setToolTip(0,name)
                mit.addChild(it)
        # objects next
        for obj in rel:
            code = self.objectslist[obj.Name]
            if (not self.form.onlyVisible.isChecked()) or obj.ViewObject.isVisible():
                it = QtGui.QTreeWidgetItem([self.labellist[obj.Name],code])
                it.setIcon(0,BimIfcElements.getIcon(obj))
                it.setToolTip(0,name)
                ok = False
                for par in obj.InListRecursive:
                    if par.Name in done:
                        if (not hasattr(par,"Hosts")) or (obj not in par.Hosts):
                            done[par.Name].addChild(it)
                            done[obj.Name] = it
                            ok = True
                            break
                if not ok:
                    self.form.treeObjects.addTopLevelItem(it)
                    done[obj.Name] = it
        self.form.treeObjects.expandAll()

    def updateDefault(self):

        import FreeCADGui,Draft
        from PySide import QtCore,QtGui
        d = self.objectslist.copy()
        d.update(self.matlist)
        for name,code in d.items():
            obj = FreeCAD.ActiveDocument.getObject(name)
            if obj:
                if (not self.form.onlyVisible.isChecked()) or obj.ViewObject.isVisible() or (Draft.getType(obj) in ["Material","MultiMaterial"]):
                    it = QtGui.QTreeWidgetItem([self.labellist[name],code])
                    it.setIcon(0,BimIfcElements.getIcon(obj))
                    it.setToolTip(0,name)
                    self.form.treeObjects.addTopLevelItem(it)
                    if obj in FreeCADGui.Selection.getSelection():
                        self.form.treeObjects.setCurrentItem(it)

    def updateClasses(self,search=""):

        from PySide import QtCore,QtGui
        self.form.treeClass.clear()

        # save as default
        FreeCAD.ParamGet("User parameter:BaseApp/Preferences/Mod/Arch").SetString("DefaultClassificationSystem",self.form.comboSystem.currentText())

        if isinstance(search,int):
            search = ""
        if self.form.search.text():
            search = self.form.search.text()
        if search:
            search = search.lower()

        system = self.form.comboSystem.currentText()
        if not system:
            return

        if not system in self.Classes:
            self.Classes[system] = self.build(system)
        if not self.Classes[system]:
            return

        for c in self.Classes[system]:
            it = None
            first = True
            if not c[1]:
                c[1] = ""
            if search:
                if (search in c[0].lower()) or (search in c[1].lower()):
                    it = QtGui.QTreeWidgetItem([c[0]+" "+c[1]])
                    it.setToolTip(0,c[1])
                    self.form.treeClass.addTopLevelItem(it)
            else:
                it = QtGui.QTreeWidgetItem([c[0]+" "+c[1]])
                it.setToolTip(0,c[1])
                self.form.treeClass.addTopLevelItem(it)
            if c[2]:
                self.addChildren(c[2],it,search)
            if it and first:
                # select first entry
                self.form.treeClass.setCurrentItem(it)
                first = False

    def addChildren(self,children,parent,search=""):

        from PySide import QtCore,QtGui
        if children:
            for c in children:
                it = None
                if not c[1]: c[1] = ""
                if search:
                    if (search in c[0].lower()) or (search in c[1].lower()):
                        it = QtGui.QTreeWidgetItem([c[0]+" "+c[1]])
                        it.setToolTip(0,c[1])
                        self.form.treeClass.addTopLevelItem(it)
                else:
                    it = QtGui.QTreeWidgetItem([c[0]+" "+c[1]])
                    it.setToolTip(0,c[1])
                    if parent:
                        parent.addChild(it)
                if c[2]:
                    self.addChildren(c[2],it,search)

    def build(self,system):
        
        # try to load the IFC first
        preset = os.path.join(FreeCAD.getUserAppDataDir(),"BIM","Classification",system+".ifc")
        if os.path.exists(preset):
            return self.build_ifc(system)
        else:
            preset = os.path.join(FreeCAD.getUserAppDataDir(),"BIM","Classification",system+".xml")
            if os.path.exists(preset):
                return self.build_xml(system)
            else:
                FreeCAD.Console.PrintError("Unable to find classification file:"+system+"\n")
                return []

    def build_ifc(self,system):
        
        # builds from ifc instead of xml

        class Item:
            def __init__(self,parent=None):
                self.parent = parent
                self.ID = None
                self.Name = None
                self.children = []

        preset = os.path.join(FreeCAD.getUserAppDataDir(),"BIM","Classification",system+".ifc")
        if not os.path.exists(preset):
            return None
        import ifcopenshell
        f = ifcopenshell.open(preset)
        classes = f.by_type("IfcClassificationReference")
        rootclass = f.by_type("IfcClassification")
        if rootclass:
            rootclass = rootclass[0]
        else:
            return None
        root = Item()
        classdict = {rootclass.id():root}
        for cl in classes:
                currentItem = Item()
                currentItem.Name = cl.Name
                currentItem.Description = cl.Description
                currentItem.ID = cl.Identification
                if cl.ReferencedSource:
                    if cl.ReferencedSource.id() in classdict:
                        currentItem.parent = classdict[cl.ReferencedSource.id()]
                        classdict[cl.ReferencedSource.id()].children.append(currentItem)
                classdict[cl.id()] = currentItem
        return [self.listize(c) for c in root.children]

    def build_xml(self,system):

        class Item:
            def __init__(self,parent=None):
                self.parent = parent
                self.ID = None
                self.Name = None
                self.children = []

        preset = os.path.join(FreeCAD.getUserAppDataDir(),"BIM","Classification",system+".xml")
        if not os.path.exists(preset):
            return None
        import codecs,re
        d = Item()
        with codecs.open(preset,"r","utf-8") as f:
            currentItem = d
            for l in f:
                if "<Item>" in l:
                    currentItem = Item(currentItem)
                    currentItem.parent.children.append(currentItem)
                if "</Item>" in l:
                    currentItem = currentItem.parent
                elif currentItem and re.findall("<ID>(.*?)</ID>",l):
                    currentItem.ID = re.findall("<ID>(.*?)</ID>",l)[0]
                elif currentItem and re.findall("<Name>(.*?)</Name>",l):
                    currentItem.Name = re.findall("<Name>(.*?)</Name>",l)[0]
                elif currentItem and re.findall("<Description>(.*?)</Description>",l) and not currentItem.Name:
                    currentItem.Name = re.findall("<Description>(.*?)</Description>",l)[0]
        return [self.listize(c) for c in d.children]


    def listize(self,item):

        return [item.ID, item.Name, [self.listize(it) for it in item.children]]

    def apply(self,item=None,col=None):

        if self.form.treeObjects.selectedItems() and len(self.form.treeClass.selectedItems()) == 1:
            c = self.form.treeClass.selectedItems()[0].text(0)
            if self.form.checkPrefix.isChecked():
                c = self.form.comboSystem.currentText() + " " + c
            for m in self.form.treeObjects.selectedItems():
                if m.toolTip(0):
                    m.setText(1,c)

    def rename(self):

        if self.form.treeObjects.selectedItems() and len(self.form.treeClass.selectedItems()) == 1:
            c = self.form.treeClass.selectedItems()[0].toolTip(0)
            for m in self.form.treeObjects.selectedItems():
                if m.toolTip(0):
                    m.setText(0,c)

    def accept(self):

        if not self.isEditing:
            changed = False
            for row in range(self.form.treeObjects.topLevelItemCount()):
                child = self.form.treeObjects.topLevelItem(row)
                items = [child]
                items.extend([child.child(childrow) for childrow in range(child.childCount())])
                for item in items:
                    code = item.text(1)
                    label = item.text(0)
                    if item.toolTip(0):
                        obj = FreeCAD.ActiveDocument.getObject(item.toolTip(0))
                        if obj:
                            if code != obj.StandardCode:
                                if not changed:
                                    FreeCAD.ActiveDocument.openTransaction("Change standard codes")
                                    changed = True
                                obj.StandardCode = code
                            if label != obj.Label:
                                if not changed:
                                    FreeCAD.ActiveDocument.openTransaction("Change standard codes")
                                    changed = True
                                obj.Label = label
            if changed:
                FreeCAD.ActiveDocument.commitTransaction()
                FreeCAD.ActiveDocument.recompute()
        else:
            code = self.form.treeClass.selectedItems()[0].text(0)
            if "StandardCode" in self.isEditing.PropertiesList:
                FreeCAD.ActiveDocument.openTransaction("Change standard codes")
                if self.form.checkPrefix.isChecked():
                    code = self.form.comboSystem.currentText() + " " + code
                self.isEditing.StandardCode = code
                if hasattr(self.isEditing.ViewObject,"Proxy") and hasattr(self.isEditing.ViewObject.Proxy,"setTaskValue"):
                    self.isEditing.ViewObject.Proxy.setTaskValue("FieldCode",code)
                FreeCAD.ActiveDocument.commitTransaction()
                FreeCAD.ActiveDocument.recompute()
        p = FreeCAD.ParamGet("User parameter:BaseApp/Preferences/Mod/BIM")
        p.SetInt("BimClassificationDialogWidth",self.form.width())
        p.SetInt("BimClassificationDialogHeight",self.form.height())
        self.form.hide()
        return True

    def onUpArrow(self):

        if self.form:
            i = self.form.treeClass.currentItem()
            if self.form.treeClass.itemAbove(i):
                self.form.treeClass.setCurrentItem(self.form.treeClass.itemAbove(i))

    def onDownArrow(self):

        if self.form:
            i = self.form.treeClass.currentItem()
            if self.form.treeClass.itemBelow(i):
                self.form.treeClass.setCurrentItem(self.form.treeClass.itemBelow(i))

